// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <ktest/ktest.h>

#include <nebula/compute/api.h>
#include <tests/compute/kernels/test_util.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/future/thread_pool.h>

namespace nebula {

using internal::ThreadPool;

namespace compute {

namespace {

void TestRandomWithOptions(int64_t length, const RandomOptions& random_options) {
  ExecBatch input({}, length);
  ASSERT_OK_AND_ASSIGN(Datum result, call_function("random", input, &random_options));
  const auto result_array = result.make_array();
  ValidateOutput(*result_array);
  ASSERT_EQ(result_array->length(), length);
  ASSERT_EQ(result_array->null_count(), 0);
  AssertTypeEqual(result_array->type(), float64());

  if (length > 0) {
    // verify E(X), E(X^2) is near theory
    double sum = 0, square_sum = 0;
    const double* values = result_array->data()->get_values<double>(1);
    for (int64_t i = 0; i < length; ++i) {
      const double value = values[i];
      ASSERT_GE(value, 0);
      ASSERT_LT(value, 1);
      sum += value;
      square_sum += value * value;
    }
    const double E_X = 0.5;
    const double E_X2 = 1.0 / 12 + E_X * E_X;
    ASSERT_NEAR(sum / length, E_X, E_X * 0.02);
    ASSERT_NEAR(square_sum / length, E_X2, E_X2 * 0.02);
  }
}

}  // namespace

TEST(TestRandom, Seed) {
  const int kCount = 100000;
  auto random_options = RandomOptions::FromSeed(/*seed=*/0);
  TestRandomWithOptions(kCount, random_options);
}

TEST(TestRandom, SystemRandom) {
  const int kCount = 100000;
  auto random_options = RandomOptions::FromSystemRandom();
  TestRandomWithOptions(kCount, random_options);
}

TEST(TestRandom, SeedIsDeterministic) {
  const int kCount = 100;
  auto random_options = RandomOptions::FromSeed(/*seed=*/0);

  ExecBatch input({}, kCount);
  ASSERT_OK_AND_ASSIGN(Datum first_call, call_function("random", input, &random_options));
  ASSERT_OK_AND_ASSIGN(Datum second_call, call_function("random", input, &random_options));
  AssertDatumsEqual(first_call, second_call);
}

TEST(TestRandom, SystemRandomDifferentResultsSingleThreaded) {
  const int kCount = 100;
  auto random_options = RandomOptions::FromSystemRandom();
  ExecBatch input({}, kCount);
  ASSERT_OK_AND_ASSIGN(Datum first_datum, call_function("random", input, &random_options));
  ASSERT_OK_AND_ASSIGN(Datum second_datum,
                       call_function("random", input, &random_options));
  ASSERT_FALSE(first_datum.equals(second_datum));
}

TEST(TestRandom, SystemRandomDifferentResultsMultiThreaded) {
  const int kCount = 100;
  const int kThreadCount = 8;
  const int kCallCount = 200;

  ASSERT_OK_AND_ASSIGN(auto pool, ThreadPool::create(kThreadCount));

  auto random_options = RandomOptions::FromSystemRandom();
  ExecBatch input({}, kCount);
  std::vector<Future<Datum>> futures;

  for (int i = 0; i < kCallCount; ++i) {
    futures.push_back(DeferNotOk(
        pool->submit([&]() { return call_function("random", input, &random_options); })));
  }
  std::vector<Datum> call_results(kCallCount);
  for (int i = 0; i < kCallCount; ++i) {
    ASSERT_OK_AND_ASSIGN(call_results[i], futures[i].result());
  }
  for (int i = 0; i < kThreadCount - 1; ++i) {
    for (int j = i + 1; j < kThreadCount; ++j) {
      ASSERT_FALSE(call_results[i].equals(call_results[j]));
    }
  }
}

TEST(TestRandom, ZeroLength) {
  auto random_options = RandomOptions::FromSystemRandom();
  TestRandomWithOptions(0, random_options);
}

}  // namespace compute
}  // namespace nebula
